import os
import numpy as np
from PIL import Image
from mindspore import ops
from mindspore.dataset import GeneratorDataset
from mindspore.dataset import vision, transforms
from mindspore.dataset.vision import Inter

'''
单标签,不再考虑四种病变的问题
'''
# root_dir:dataset\train
class DDRDataset:
    def __init__(self, root_dir):
        self.root_dir = root_dir
        self.labels=os.listdir(os.path.join(root_dir,'label'))
        self.images = os.listdir(os.path.join(root_dir, 'image'))

    def __getitem__(self, i):
        imgName = self.images[i]
        imagePath = os.path.join(self.root_dir, 'image', imgName)
        image = Image.open(imagePath)

        labelName=self.labels[i]
        labelPath=os.path.join(self.root_dir,'label',labelName)
        label=Image.open(labelPath)

        return image, label

    def __len__(self):
        return len(self.images)


def generator_func(dataset):
    for i in range(len(dataset)):
        yield dataset[i]

path_to_train_dir = '../dataset_test/train'
path_to_val_dir = '../dataset_test/validate'
path_to_test_dir = '../dataset_test/test'


# Create dataset
train_dataset = DDRDataset(path_to_train_dir)
val_dataset = DDRDataset(path_to_val_dir)
test_dataset = DDRDataset(path_to_test_dir)

# Create dataloader
train_data = GeneratorDataset(source=lambda:generator_func(train_dataset), column_names=['image', 'label'], shuffle=True)
val_data = GeneratorDataset(source=lambda:generator_func(val_dataset), column_names=['image', 'label'], shuffle=False)
test_data = GeneratorDataset(source=lambda:generator_func(test_dataset), column_names=['image', 'label'], shuffle=False)

'''
transform_param 
{ 
    #训练集预处理需要进行的操作
    mirror: true #以50%的概率水平翻转图片
    mean_value: 13.1248 #RGB3通道需要减去的均值
    mean_value: 36.4886
    mean_value: 74.6864
}
'''

imgTransfrom=[
    vision.Resize([712,1072],Inter.BILINEAR),
    vision.RandomHorizontalFlip(),#以50%的概率随机翻转图像
    vision.Normalize(mean=(123.675,116.28,103.53),std=(58.395,57.12,57.375)),
    vision.HWC2CHW()
]

#标签需要转换为张量,便于损失函数计算前景区域和背景区域
labelTransfrom=[
    vision.Resize([712,1072],Inter.BILINEAR),
    vision.ToTensor()
]
train_data=train_data.map(input_columns=["image"], operations=imgTransfrom)
train_data=train_data.map(input_columns=["label"], operations=labelTransfrom)
train_data=train_data.batch(batch_size=1)

val_data=val_data.map(input_columns=["image"], operations=imgTransfrom)
val_data=val_data.map(input_columns=["label"], operations=labelTransfrom)
val_data=val_data.batch(batch_size=1)

test_data=test_data.map(input_columns=["image"], operations=imgTransfrom)
test_data=test_data.map(input_columns=["label"], operations=labelTransfrom)
test_data=test_data.batch(batch_size=1)